#include "CudaGLComputeApp.hpp"
#include "cudaPostProcess.cuh"


//#define USE_GL_COMPUTE


void CudaGLComputeApp::updateFrame()
{
  m_camera->update();
  
  // --------
  
  if (aer::EventHandler::IsKeyPressed( sf::Keyboard::Escape )) quit();
  if (aer::EventHandler::IsKeyDown( sf::Keyboard::T ))  m_bDisplayStats = !m_bDisplayStats;
  
  bool bKernelRadiusUpdated = false;
  
  if (aer::EventHandler::IsKeyPressed( sf::Keyboard::Add )) {
    m_kernelRadius = std::min( aer::I32(m_kernelRadius+1), aer::I32(kMaxKernelRadius));
    bKernelRadiusUpdated = true;
  }  
  if (aer::EventHandler::IsKeyPressed( sf::Keyboard::Subtract )) {
    m_kernelRadius = std::max( aer::I32(m_kernelRadius-1), aer::I32(0));
    bKernelRadiusUpdated = true;
  }
  
  if (aer::EventHandler::IsKeyPressed( sf::Keyboard::Space )) {
    m_bUseCUDA = !m_bUseCUDA;

    printf("Kernel : %s\n", (m_bUseCUDA)?"CUDA":"GLSL Compute Shader");
  }
  
  // --------

  if (bKernelRadiusUpdated) {
    printf("Kernel radius : %d\n", m_kernelRadius);
  }
}

void CudaGLComputeApp::drawFrame()
{
  m_firstPass.fbo.begin();
  renderScene();
  m_firstPass.fbo.end();

  renderPostProcess();
}

void CudaGLComputeApp::init()
{
  /// Init the camera
  aer::View view( aer::Vector3( 0.0f, 0.0f, 30.0f),
                  aer::Vector3( 0.0f, 0.0f, 0.0f),
                  aer::Vector3( 0.0f, 1.0f, 0.0f) );
 
  aer::Frustum frustum( 60.0f, 1.0f, 0.1f, 500.0f);

  m_camera = new aer::FreeCamera( view, frustum);
  m_camera->setMoveCoefficient( 1.5f );


  /// OpenGL settings
  glClearColor( 0.5f, 0.15f, 0.15f, 1.0f);
  glEnable( GL_DEPTH_TEST );

  glGenQueries( 1, &m_interop.query);//

  initGeometry();
  initTexture();
  initShader();
}

void CudaGLComputeApp::deinit()
{
  AER_SAFE_DELETE( m_camera );

  cudaGraphicsUnregisterResource( m_interop.texSRC_CUDAResource );
  cudaGraphicsUnregisterResource( m_interop.texDST_CUDAResource );

  glDeleteQueries( 1, &m_interop.query );

  // TODO
}

void CudaGLComputeApp::initGeometry()
{  
  m_sceneObject.init( 4.0f, 32u);

  CHECKGLERROR();
}

void CudaGLComputeApp::initTexture()
{ 
// -- FirstPass buffer textures (RGBA + DEPTH)

  m_firstPass.texRGBA.generate();
  m_firstPass.texRGBA.bind();
  m_firstPass.texRGBA.setFilters( aer::Texture::NEAREST, aer::Texture::NEAREST);
  m_firstPass.texRGBA.setWraps( aer::Texture::CLAMP_TO_EDGE, aer::Texture::CLAMP_TO_EDGE);
  m_firstPass.texRGBA.resize( kSize, kSize, GL_RGBA8);

  //m_firstPass.texRGBA.setImage(aer::Texture::TEXTURE_2D, GL_RGBA16F, kSize, kSize, GL_RGBA, GL_UNSIGNED_BYTE, NULL);

  // Register the texture by CUDA to be later bound as a read-only CUDA texture
  CHECK_CUDA(
  cudaGraphicsGLRegisterImage(&m_interop.texSRC_CUDAResource, m_firstPass.texRGBA.getId(),
                              GL_TEXTURE_2D, cudaGraphicsRegisterFlagsReadOnly));


  m_firstPass.texDEPTH.generate();
  m_firstPass.texDEPTH.bind();
  m_firstPass.texDEPTH.setFilters( aer::Texture::NEAREST, aer::Texture::NEAREST);
  m_firstPass.texDEPTH.setWraps( aer::Texture::CLAMP_TO_EDGE, aer::Texture::CLAMP_TO_EDGE);
  m_firstPass.texDEPTH.resize( kSize, kSize, GL_DEPTH_COMPONENT24);

  aer::Texture::Unbind( aer::Texture::TEXTURE_2D );


// -- FirstPass FBO

  m_firstPass.fbo.generate();
  m_firstPass.fbo.bind();
  {
    aer::FrameBuffer &fbo = m_firstPass.fbo;
    
    fbo.attachColorBuffer(   &m_firstPass.texRGBA,  GL_COLOR_ATTACHMENT0);
    fbo.attachSpecialBuffer( &m_firstPass.texDEPTH, GL_DEPTH_ATTACHMENT);
    
    AER_CHECK( aer::FrameBuffer::CheckStatus() );
  }
  m_firstPass.fbo.unbind();
  m_firstPass.fbo.resize( kSize, kSize);


// -- Kernel's output Texture

  m_interop.texDST.generate();
  m_interop.texDST.bind();
  m_interop.texDST.setFilters( aer::Texture::NEAREST, aer::Texture::NEAREST);
  m_interop.texDST.setWraps( aer::Texture::CLAMP_TO_EDGE, aer::Texture::CLAMP_TO_EDGE);
  m_interop.texDST.resize( kSize, kSize, GL_RGBA8);

  aer::Texture::Unbind( aer::Texture::TEXTURE_2D );

  // Register the texture by CUDA to be later bound to a CUDA Surface reference
  CHECK_CUDA(
  cudaGraphicsGLRegisterImage(&m_interop.texDST_CUDAResource, m_interop.texDST.getId(),
                              GL_TEXTURE_2D, cudaGraphicsRegisterFlagsSurfaceLoadStore));
  
  CHECKGLERROR();
}

void CudaGLComputeApp::initShader()
{
  glswInit();  

  const char *kShadersPath = "../CudaGLCompute/data/"; 
  glswSetPath( kShadersPath, ".glsl");

  glswAddDirectiveToken("*", "#version 420 core"); 
  
  // GL_ARB_compute_shader is core in GL 4.3
  glswAddDirectiveToken("CS", "#extension GL_ARB_compute_shader : enable");
  
  // Set static Kernel parameters
  char glslToken[64];
  sprintf( glslToken, "#define BLOCK_DIM %d", kBlockDim);
  glswAddDirectiveToken( "CS", glslToken);
  sprintf( glslToken, "#define MAX_RADIUS %d", kMaxKernelRadius);
  glswAddDirectiveToken("CS", glslToken);

  
  // Scene shader
  m_scenePgm.generate();
    m_scenePgm.addShader( aer::Program::VERTEX_SHADER,   "PassThrough.VS");
    m_scenePgm.addShader( aer::Program::FRAGMENT_SHADER, "PassThrough.FS");
  AER_CHECK( m_scenePgm.link() );

  // Compute Shader kernel to blur input image
  m_computePgm.generate();
  AER_CHECK( m_computePgm.create( aer::Program::COMPUTE_SHADER, "PostProcess.CS") );

  // Map a texture to the screen
  m_screenPgm.generate();
    m_screenPgm.addShaderSrc( aer::Program::VERTEX_SHADER,   
                              aer::ScreenQuad::GetVertexShader());
    m_screenPgm.addShaderSrc( aer::Program::FRAGMENT_SHADER, 
                              aer::ScreenQuad::GetFragmentShader());
  AER_CHECK( m_screenPgm.link() );


  glswShutdown();

  CHECKGLERROR();
}


void CudaGLComputeApp::renderScene()
{
  glClear( GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT );

  m_scenePgm.bind();
  {
    const aer::Matrix4x4 &mvp = m_camera->getViewProjectionMatrix();
    m_scenePgm.setUniform( "uModelViewProjMatrix", mvp);
    m_sceneObject.draw();
  }
  m_scenePgm.unbind();

  CHECKGLERROR();
}

void CudaGLComputeApp::renderPostProcess()
{
// -- Apply postProcess in a GPU kernel

  if (m_bUseCUDA) {
    launchCUDAKernel();    
  } else {
    launchGLCompute();  
  }


// -- Render output image

  glDisable( GL_DEPTH_TEST );  
  m_screenPgm.bind();
  {
    aer::Texture &tex = m_interop.texDST;
    tex.bind(0);

    m_screenPgm.setUniform( "uTex", 0);
    aer::ScreenQuad::Draw();
    
    tex.unbind();
  }
  m_screenPgm.unbind();
  glEnable( GL_DEPTH_TEST );

  CHECKGLERROR();
}


void CudaGLComputeApp::launchGLCompute()
{
  glBeginQuery( GL_TIME_ELAPSED, m_interop.query);
  m_computePgm.bind();
    // Bind source texture in TEXTURE_UNIT 0
    m_firstPass.texRGBA.bind(0);
    m_computePgm.setUniform( "uSrcTex", 0);

    // Bind destination texture to IMAGE_UNIT 0
    glBindImageTexture( 0, m_interop.texDST.getId(), 0, GL_FALSE, 0, 
                        GL_WRITE_ONLY, GL_RGBA8);
    m_computePgm.setUniform( "uDstImg", 0);

    m_computePgm.setUniform( "uRadius", aer::I32(m_kernelRadius));

    // Launch kernel
    glDispatchCompute( kSize/kBlockDim, kSize/kBlockDim, 1);

    // Unbind src & dst textures
    m_firstPass.texRGBA.unbind();
    //glBindImageTextures( 0, 1, NULL); // GL 4.4+
  m_computePgm.unbind();
  glEndQuery( GL_TIME_ELAPSED );


  if (m_bDisplayStats) {
    GLint nanosecond = 0;
    glGetQueryObjectiv( m_interop.query, GL_QUERY_RESULT, &nanosecond);
    float ms = nanosecond/1000000.0f;
    fprintf(stderr, "%.3f ms\n", ms);
    m_bDisplayStats = false;
  }

  CHECKGLERROR();
}

void CudaGLComputeApp::launchCUDAKernel()
{
  m_interop.cudaTimer.Start();

  cudaArray *d_inArray = NULL;
  CHECK_CUDA(cudaGraphicsMapResources( 1, &m_interop.texSRC_CUDAResource));
  CHECK_CUDA(cudaGraphicsSubResourceGetMappedArray( &d_inArray, m_interop.texSRC_CUDAResource, 0, 0));

  cudaArray *d_outArray = NULL;
  CHECK_CUDA(cudaGraphicsMapResources( 1, &m_interop.texDST_CUDAResource));
  CHECK_CUDA(cudaGraphicsSubResourceGetMappedArray( &d_outArray, m_interop.texDST_CUDAResource, 0, 0));

  
  const dim3 gridDim( kSize/kBlockDim, kSize/kBlockDim);
  const dim3 blockDim( kBlockDim, kBlockDim);
  const size_t tileWidth = (kBlockDim + 2*m_kernelRadius);
  const size_t smemSize = tileWidth * tileWidth * 4 * sizeof(unsigned char);

  launch_cuda_kernel( gridDim, blockDim, smemSize, 
                      d_inArray, d_outArray, kSize, tileWidth, m_kernelRadius);

  
  CHECK_CUDA(cudaGraphicsUnmapResources( 1, &m_interop.texSRC_CUDAResource));
  CHECK_CUDA(cudaGraphicsUnmapResources( 1, &m_interop.texDST_CUDAResource));

  m_interop.cudaTimer.Stop();

  // Note : CUDATimer is ineffective when screen vertical sync is enabled.
  if (m_bDisplayStats) {
    float ms = m_interop.cudaTimer.Elapsed();
    fprintf(stderr, "%.3f ms\n", ms);
    m_bDisplayStats = false;
  }

  CHECKCUDAERROR();
}
